# -*- coding: utf-8 -*-
# @Time    : 2025/9/19 11:11
# @Author  : chenmh
# @File    : utils.py
# @Desc: 工具类

import torch, logging, os, datetime
from typing import Tuple
import numpy as np
from sklearn.metrics import accuracy_score, roc_auc_score, recall_score, precision_score, classification_report, \
    f1_score
import logging, random
from typing import Union
from pathlib import Path


def calc_metrics(y_true: Union[np.ndarray, torch.Tensor], y_pred: Union[np.ndarray, torch.Tensor],
                 y_prob: Union[np.ndarray, torch.Tensor], logger: logging.Logger) -> float:
    if isinstance(y_true, torch.Tensor):
        y_true = y_true.detach().cpu().numpy()
        y_pred = y_pred.detach().cpu().numpy()
        y_prob = y_prob.detach().cpu().numpy()
    auc = roc_auc_score(y_true=y_true, y_score=y_prob)
    logger.info(str(classification_report(y_true=y_true, y_pred=y_pred, zero_division=0)))
    logger.info(f"AUC:{auc}")
    return -auc


def set_seed(seed=42):
    random.seed(seed)  # Python
    np.random.seed(seed)  # NumPy
    torch.manual_seed(seed)  # CPU
    torch.cuda.manual_seed(seed)  # GPU
    torch.cuda.manual_seed_all(seed)  # GPU
    torch.backends.cudnn.deterministic = True  #
    torch.backends.cudnn.benchmark = False  #


def tensor_difference(tensor: torch.Tensor) -> torch.Tensor:
    n, m = tensor.shape
    # 计算从第二个元素开始到最后一个元素与前一个元素的差
    diffs = tensor[:, 1:] - tensor[:, :-1]
    # 在前面补0，保持与原tensor相同的形状
    result = torch.cat([torch.zeros(n, 1, device=tensor.device), diffs], dim=1)
    return result


def setup_logging(save_dir: str, log_file_name: str, log_time: str) -> Tuple[logging.Logger, str]:
    """配置日志：同时输出到控制台和文件"""
    # 创建日志目录
    log_dir = os.path.join(save_dir, "logs")
    os.makedirs(log_dir, exist_ok=True)

    # 日志文件名（包含时间戳，避免覆盖）
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    if log_time: timestamp = log_time
    log_file = os.path.join(log_dir, f"{log_file_name}_{timestamp}.log")
    # 配置日志基本设置
    logging.basicConfig(
        level=logging.INFO,  # 日志级别：INFO及以上
        format="%(asctime)s - %(levelname)s - %(message)s",  # 日志格式
        handlers=[
            logging.FileHandler(log_file),  # 写入文件
            logging.StreamHandler()  # 输出到控制台
        ]
    )

    # 返回日志器
    logger = logging.getLogger(__name__)
    logger.info(f"日志已配置，文件保存路径：{log_file}")
    return logger, str(timestamp)


def get_latest_pth_file(folder: str = './logs'):
    folder = Path(folder)
    if not folder.exists():
        raise FileNotFoundError(f"文件夹 {folder} 不存在")
    pth_files = list(folder.glob('*.pth'))
    if not pth_files:
        raise FileNotFoundError(f"在 {folder} 中未找到 .pth 文件")
    # 按修改时间降序排列，取最新的
    latest_file = max(pth_files, key=lambda f: f.stat().st_mtime)
    model_full_path = str(latest_file).split("\\")[-1]
    return model_full_path, "_".join(model_full_path.split(".")[0].split("_")[-2:])
